#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Create a video from the visualization PNG sequence.
"""

import os
import sys
import glob
import argparse
import numpy as np

# Try to import necessary libraries
try:
    import imageio
    IMAGEIO_AVAILABLE = True
except ImportError:
    IMAGEIO_AVAILABLE = False
    print("Warning: imageio not installed. Install with: pip install imageio[ffmpeg]")

try:
    import cv2
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    print("Warning: OpenCV not installed. Install with: pip install opencv-python")

def create_video_with_imageio(image_folder, output_path, fps=10):
    """Create video using imageio library."""
    if not IMAGEIO_AVAILABLE:
        print("Error: imageio is not available")
        return False
    
    # Get all PNG files sorted by frame number
    images = sorted(glob.glob(os.path.join(image_folder, "scene_*.png")))
    
    if not images:
        print(f"No scene_*.png files found in {image_folder}")
        return False
    
    print(f"Found {len(images)} frames")
    print(f"Creating video at {fps} fps...")
    
    # Create video writer
    try:
        with imageio.get_writer(output_path, fps=fps, codec='libx264', quality=8) as writer:
            for i, image_path in enumerate(images):
                img = imageio.imread(image_path)
                writer.append_data(img)
                
                # Progress indicator
                if (i + 1) % 10 == 0:
                    print(f"  Processed {i + 1}/{len(images)} frames")
        
        print(f"✓ Video saved to: {output_path}")
        return True
        
    except Exception as e:
        print(f"Error creating video with imageio: {e}")
        return False

def create_video_with_opencv(image_folder, output_path, fps=10):
    """Create video using OpenCV."""
    if not CV2_AVAILABLE:
        print("Error: OpenCV is not available")
        return False
    
    # Get all PNG files sorted by frame number
    images = sorted(glob.glob(os.path.join(image_folder, "scene_*.png")))
    
    if not images:
        print(f"No scene_*.png files found in {image_folder}")
        return False
    
    print(f"Found {len(images)} frames")
    
    # Read first image to get dimensions
    frame = cv2.imread(images[0])
    if frame is None:
        print(f"Error: Could not read {images[0]}")
        return False
    
    height, width, layers = frame.shape
    print(f"Frame size: {width}x{height}")
    print(f"Creating video at {fps} fps...")
    
    # Define codec and create VideoWriter
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can also try 'XVID'
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    try:
        for i, image_path in enumerate(images):
            frame = cv2.imread(image_path)
            if frame is None:
                print(f"Warning: Could not read {image_path}, skipping")
                continue
            
            out.write(frame)
            
            # Progress indicator
            if (i + 1) % 10 == 0:
                print(f"  Processed {i + 1}/{len(images)} frames")
        
        out.release()
        cv2.destroyAllWindows()
        
        print(f"✓ Video saved to: {output_path}")
        return True
        
    except Exception as e:
        print(f"Error creating video with OpenCV: {e}")
        return False

def create_ffmpeg_command(image_folder, output_path, fps=10):
    """Generate FFmpeg command for video creation."""
    pattern = os.path.join(image_folder, "scene_%05d.png")
    
    # Basic command
    cmd = f"ffmpeg -y -framerate {fps} -i '{pattern}' -c:v libx264 -pix_fmt yuv420p -crf 23 '{output_path}'"
    
    # Alternative with more compatibility
    cmd_alt = f"ffmpeg -y -pattern_type glob -framerate {fps} -i '{os.path.join(image_folder, 'scene_*.png')}' -c:v libx264 -pix_fmt yuv420p -crf 23 '{output_path}'"
    
    print("FFmpeg commands you can use:")
    print(f"\nOption 1 (numbered sequence):")
    print(cmd)
    print(f"\nOption 2 (glob pattern):")
    print(cmd_alt)
    
    return cmd

def main():
    parser = argparse.ArgumentParser(description='Create video from visualization PNG sequence')
    parser.add_argument('--input', default='../../output/visualization', help='Input folder with PNG files')
    parser.add_argument('--output', default='dragonfly_simulation.mp4', help='Output video file')
    parser.add_argument('--fps', type=int, default=10, help='Frames per second (default: 10)')
    parser.add_argument('--method', choices=['auto', 'imageio', 'opencv', 'ffmpeg'], 
                       default='auto', help='Method to use for video creation')
    
    args = parser.parse_args()
    
    # Get absolute paths
    if not os.path.isabs(args.input):
        args.input = os.path.abspath(args.input)
    if not os.path.isabs(args.output):
        args.output = os.path.abspath(args.output)
    
    print("="*60)
    print("Dragonfly Simulation - Video Creator")
    print("="*60)
    print(f"Input folder: {args.input}")
    print(f"Output video: {args.output}")
    print(f"Frame rate: {args.fps} fps")
    print("="*60)
    
    if not os.path.exists(args.input):
        print(f"Error: Input folder '{args.input}' does not exist")
        return 1
    
    success = False
    
    if args.method == 'auto':
        # Try imageio first, then OpenCV
        if IMAGEIO_AVAILABLE:
            print("\nUsing imageio...")
            success = create_video_with_imageio(args.input, args.output, args.fps)
        
        if not success and CV2_AVAILABLE:
            print("\nTrying OpenCV...")
            output_cv2 = args.output.replace('.mp4', '_opencv.mp4')
            success = create_video_with_opencv(args.input, output_cv2, args.fps)
        
        if not success:
            print("\nPython libraries not available. Use FFmpeg directly:")
            create_ffmpeg_command(args.input, args.output, args.fps)
    
    elif args.method == 'imageio':
        success = create_video_with_imageio(args.input, args.output, args.fps)
    
    elif args.method == 'opencv':
        success = create_video_with_opencv(args.input, args.output, args.fps)
    
    elif args.method == 'ffmpeg':
        create_ffmpeg_command(args.input, args.output, args.fps)
        print("\nRun the command above in your terminal to create the video.")
        return 0
    
    if success:
        print("\n✓ Video creation successful!")
        # Check if the file was created with a different name (e.g., _opencv suffix)
        if not os.path.exists(args.output):
            # Look for alternative output files
            base_name = os.path.splitext(args.output)[0]
            alternatives = [base_name + "_opencv.mp4", base_name + "_imageio.mp4"]
            for alt in alternatives:
                if os.path.exists(alt):
                    print(f"  View your video: {alt}")
                    break
        else:
            print(f"  View your video: {args.output}")
        
        # Calculate video duration
        images = sorted(glob.glob(os.path.join(args.input, "scene_*.png")))
        duration = len(images) / args.fps
        print(f"  Duration: {duration:.1f} seconds")
        print(f"  Total frames: {len(images)}")
    else:
        print("\n✗ Video creation failed!")
        print("To install missing dependencies:")
        print("  pip install imageio[ffmpeg]  # For imageio with FFmpeg")
        print("  pip install opencv-python    # For OpenCV")
        print("Or use FFmpeg directly with the commands shown above.")
    
    return 0 if success else 1

if __name__ == "__main__":
    sys.exit(main()) 